#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
机器人前向动力学程序 - 使用 PyBullet 滑块控制
改动：将“摄像机偏航 (yaw)”滑条放在最上方
"""

import os
import time
import pybullet as p
import pybullet_data

# ======== 摄像机默认参数 ========
CAM_DISTANCE = 0.5
CAM_YAW_INIT = 50          # 初始偏航角（滑条初值）
CAM_PITCH    = -30
CAM_TARGET   = [0, 0, 0.5]

# 键值（仅保留 ESC 退出）
KEY_ESC = 65307

# URDF 文件路径
urdf_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "robot2.urdf")


class RobotController:
    """机器人控制器：模型加载、滑条创建、关节与摄像机控制"""

    def __init__(self, urdf_path):
        if not os.path.exists(urdf_path):
            raise FileNotFoundError(f"找不到 URDF 文件: {urdf_path}")
        self.urdf_path = urdf_path
        self.robot_id = None
        self.joint_ids = []
        self.joint_names = []
        self.joint_angles = []
        self.joint_sliders = []
        self.cam_yaw_slider = None

    # ---------- 初始化仿真 ----------
    def init_pybullet(self):
        try:
            p.connect(p.GUI)
        except p.error:
            p.connect(p.DIRECT)
            print("GUI 连接失败，已切换到 DIRECT 模式")

        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.8)
        plane_id = p.loadURDF("plane.urdf")
        p.changeDynamics(plane_id, -1, lateralFriction=1)

        start_pos = [0, 0, 0.5]
        start_ori = p.getQuaternionFromEuler([0, 0, 0])
        self.robot_id = p.loadURDF(self.urdf_path, start_pos, start_ori, useFixedBase=True)

        # 收集可控关节
        for i in range(p.getNumJoints(self.robot_id)):
            jt_type = p.getJointInfo(self.robot_id, i)[2]
            if jt_type != p.JOINT_FIXED:
                self.joint_ids.append(i)
                name = p.getJointInfo(self.robot_id, i)[1].decode("utf-8")
                self.joint_names.append(" " + name)
        self.joint_angles = [0.0] * len(self.joint_ids)
        p.setTimeStep(1 / 240)

        # 初次设置摄像机
        p.resetDebugVisualizerCamera(
            cameraDistance=CAM_DISTANCE,
            cameraYaw=CAM_YAW_INIT,
            cameraPitch=CAM_PITCH,
            cameraTargetPosition=CAM_TARGET,
        )

    # ---------- 创建滑条 ----------
    def create_sliders(self):
        # 1. 先创建“摄像机偏航”滑条 —— 会排在最顶
        self.cam_yaw_slider = p.addUserDebugParameter(
            "(yaw)", -180, 180, CAM_YAW_INIT
        )

        # 2. 再创建关节滑条
        for jn in self.joint_names:
            sid = p.addUserDebugParameter(jn, -3.14, 3.14, 0.0)
            self.joint_sliders.append(sid)

        # 3. 最后创建“重置所有关节”按钮（排在最底）
        self.reset_button = p.addUserDebugParameter("重置所有关节", 1, 0, 1)

    # ---------- 驱动关节 ----------
    def apply_joint_angles(self):
        for i, jid in enumerate(self.joint_ids):
            p.setJointMotorControl2(
                self.robot_id, jid, p.POSITION_CONTROL,
                targetPosition=self.joint_angles[i], force=500
            )

    def reset_joint_angles(self):
        self.joint_angles = [0.0] * len(self.joint_ids)
        self.apply_joint_angles()


# ============ 主程序 ============
def main():
    running = True
    rc = RobotController(urdf_file)
    rc.init_pybullet()
    rc.create_sliders()

    last_reset_val = p.readUserDebugParameter(rc.reset_button)

    try:
        while running:
            keys = p.getKeyboardEvents()
            if KEY_ESC in keys and (keys[KEY_ESC] & p.KEY_WAS_TRIGGERED):
                break

            # --- 摄像机偏航 ---
            cam_yaw = p.readUserDebugParameter(rc.cam_yaw_slider)
            p.resetDebugVisualizerCamera(
                cameraDistance=CAM_DISTANCE,
                cameraYaw=cam_yaw,
                cameraPitch=CAM_PITCH,
                cameraTargetPosition=CAM_TARGET,
            )

            # --- “重置所有关节”按钮 ---
            reset_val = p.readUserDebugParameter(rc.reset_button)
            if reset_val != last_reset_val:
                rc.reset_joint_angles()
                last_reset_val = reset_val

            # --- 读取关节滑条并驱动关节 ---
            for i, sid in enumerate(rc.joint_sliders):
                rc.joint_angles[i] = p.readUserDebugParameter(sid)
            rc.apply_joint_angles()

            p.stepSimulation()
            time.sleep(1 / 240)

    finally:
        if p.isConnected():
            p.disconnect()


if __name__ == "__main__":
    main()
